"""
统计套利策略(配对交易)
策略逻辑：
1. 选择高度相关的资产对
2. 计算价差和标准差
3. 当价差偏离均值时建仓
4. 价差回归时平仓
"""

import numpy as np
import pandas as pd
from statsmodels.tsa.statespace.tools import cfa

class PairsTradingStrategy:
    def __init__(self, symbol1, symbol2, lookback=30, z_threshold=2):
        self.symbol1 = symbol1
        self.symbol2 = symbol2
        self.lookback = lookback
        self.z_threshold = z_threshold
        self.spread_history = []
        self.position = 0
        
    def calculate_spread(self, price1, price2):
        """计算标准化价差"""
        hedge_ratio = self.calculate_hedge_ratio(price1, price2)
        spread = price1 - hedge_ratio * price2
        return (spread - np.mean(spread)) / np.std(spread)
    
    def calculate_hedge_ratio(self, price1, price2):
        """使用协整关系计算对冲比例"""
        if len(price1) < 2 or len(price2) < 2:
            return 1.0
        return np.cov(price1, price2)[0,1] / np.var(price2)
    
    def generate_signal(self, price1, price2):
        """生成交易信号"""
        if len(price1) < self.lookback or len(price2) < self.lookback:
            return 0
            
        # 计算当前价差
        current_spread = self.calculate_spread(
            price1[-self.lookback:],
            price2[-self.lookback:]
        )[-1]
        
        # 生成信号
        if current_spread > self.z_threshold and self.position <= 0:
            return -1  # 做空价差(买入symbol2，卖出symbol1)
        elif current_spread < -self.z_threshold and self.position >= 0:
            return 1   # 做多价差(买入symbol1，卖出symbol2)
        elif abs(current_spread) < 0.5 and self.position != 0:
            return 0    # 平仓
        else:
            return None # 保持现有仓位
    
    def update_position(self, signal, price1, price2):
        """更新仓位"""
        if signal == 1:
            self.position = 1
            print(f"建仓多头: {self.symbol1}/{self.symbol2}")
        elif signal == -1:
            self.position = -1
            print(f"建仓空头: {self.symbol1}/{self.symbol2}")
        elif signal == 0:
            print(f"平仓: {self.symbol1}/{self.symbol2}")
            self.position = 0

if __name__ == '__main__':
    # 示例用法
    strategy = PairsTradingStrategy("AAPL", "MSFT")
    
    # 模拟价格数据
    np.random.seed(42)
    prices1 = np.cumsum(np.random.randn(100)) + 100
    prices2 = prices1 + np.random.randn(100)*2
    
    # 模拟运行
    for i in range(30, 100):
        signal = strategy.generate_signal(prices1[:i], prices2[:i])
        if signal is not None:
            strategy.update_position(signal, prices1[i], prices2[i])